from sklearn.preprocessing import LabelEncoder
import pandas as pd


def encode(df, cols):
    for col in cols:
        le = LabelEncoder()
        df[col] = le.fit_transform(df[col])
    return df


if __name__ == '__main__':
    data = pd.read_csv('../../data/raw/train.csv')
    cols = ["BusinessTravel", "Department", "Education", "EducationField", "EnvironmentSatisfaction",
            "Gender", "JobInvolvement", "JobLevel", "JobRole",
            "JobSatisfaction", "MaritalStatus", "OverTime", "PerformanceRating",
            "RelationshipSatisfaction", "TrainingTimesLastYear", "StockOptionLevel", "WorkLifeBalance"]
    data = encode(data, cols)
    print(data.head())